package org.demos.springaiargdemo1.service;

import cn.hutool.core.collection.CollUtil;
import org.demos.springaiargdemo1.model.dto.RAGResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

@Service
public class RAGService {

    private static final Logger log = LoggerFactory.getLogger(DocumentService.class);

    private final VectorStore vectorStore;
    private final ChatClient chatClient;

    // 检索的相关文档数量
    private static final int TOP_K = 5;

    public RAGService(VectorStore vectorStore, ChatClient.Builder builder) {
        this.vectorStore = vectorStore;
        this.chatClient = builder.build();
    }

    /**
     * 处理RAG问答 - 完全响应式实现
     */
    public Flux<RAGResponse> generateResponse(String question) {
        // 1. 从向量存储中检索相关文档 - 使用响应式方式
        SearchRequest searchRequest = SearchRequest.builder()
                .query(question)
                .topK(TOP_K)
                .build();
        
        return Mono.fromCallable(() -> {
            // 在单独的线程中执行阻塞操作
            return vectorStore.similaritySearch(searchRequest);
        })
        .subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()) // 使用弹性线程池
        .flatMapMany(relevantDocs -> {
            // 4. 创建消息列表
            List<Message> messages = new ArrayList<>();
            List<Document> filteredDocs = relevantDocs.stream().filter(doc -> Objects.nonNull(doc.getScore()) &&doc.getScore() > 0.5).toList();
            log.info("检索到 {} 个相关文档", filteredDocs.size());
            // 2. 构建系统提示信息
            String context = buildContext(filteredDocs);
            SystemMessage systemMessage = getSystemMessage(context);
            messages.add(systemMessage);

            
            // 3. 构建用户问题
            UserMessage userMessage = new UserMessage(question);
            messages.add(userMessage);
            List<RAGResponse.RetrievedDocument> docs = filteredDocs.stream().map(doc ->
                 RAGResponse.RetrievedDocument.builder()
                        .content(doc.getText())
                         .title(String.valueOf(doc.getMetadata().get("lang")))
                         .score(doc.getScore())
                         .id(doc.getId())
                        .build()
            ).toList();
            // 5. 使用ChatClient生成回答
            return chatClient.prompt(new Prompt(messages)).stream().content().map(content ->
                 RAGResponse.builder()
                        .content(content)
                        .documents(docs)
                        .isLast(true)
                        .build()
            );
        })
        .onErrorResume(throwable -> {
            log.error("RAG处理失败: {}", throwable.getMessage(), throwable);
            // 降级到普通聊天模式
            log.info("降级到普通聊天模式");
            UserMessage userMessage = new UserMessage(question);
            return chatClient.prompt(new Prompt(List.of(userMessage))).stream().content()
                    .map(content ->
                         RAGResponse.builder()
                                .content(content)
                                .isLast(true)
                                .build()
                    );
        });
    }

    /**
     * 构建上下文信息
     */
    private String buildContext(List<Document> documents) {
        if (CollUtil.isEmpty(documents)){
            return "没有检索到有用信息。";
        }
        StringBuilder context = new StringBuilder();
        context.append("以下是一些相关的调试信息：\n\n");

        for (Document doc : documents) {
            context.append("---\n");
            context.append(doc.getText());
            context.append("\n");
        }

        return context.toString();
    }

    /**
     * 获取系统消息
     */
    private SystemMessage getSystemMessage(String context) {
        return new SystemMessage("""
            你是一个专业的调试助手。我会给你一个问题以及检索到的相关资料(可能没有)。
            请基于提供的上下文信息来回答问题。如果上下文信息不足以回答问题，请明确说明。
            
            上下文信息：
            %s
            
            请记住：
            1. 保持回答准确、专业
            2. 如果需要代码示例，请提供清晰的注释
            3. 如果有多个可能的解决方案，请列出优缺点
            4. 如果上下文信息不足，请说明还需要什么额外信息
            """.formatted(context));
    }
}
